# All content Copyright 2010 Cyrus Omar <cyrus.omar@gmail.com> unless otherwise
# specified.
#
# Contributors:
#     Cyrus Omar <cyrus.omar@gmail.com>
#
# This file is part of, and licensed under the terms of, the atomic-hedgehog
# package.
#
# The atomic-hedgehog package is free software: you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 3 of the License,
# or (at your option) any later version.
#
# The atomic-hedgehog package is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
# License for more details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with the atomic-hedgehog package. If not, see <http://www.gnu.org/licenses/>.
"""This is a sample script used for testing."""

def get_device():
    from ahh.util.opencl import SimpleDeviceContext
    global ctx
    ctx = SimpleDeviceContext.get_somehow()

def init_sim():
    from ahh.sim import Simulation
    global sim
    sim = Simulation(ctx)

def init_neurons():
    from ahh.sim.models import ReducedLIF
    global N, NE, NI
    global neurons

    N = 4000
    NE = N * 4/5
    NI = N - NE

    neurons = ReducedLIF(sim, N, "LIF")
    neurons.tau = 20.0
    neurons.v_reset = 0.0
    neurons.v_thresh = 20.0
    neurons.abs_refractory_period = 5.0

def init_synapses():
    from ahh.sim.inputs import ExpSynapse
    global ge, gi

    ge = ExpSynapse(neurons, "ge")
    ge.tau = 5.0
    ge.reversal = 60.0

    gi = ExpSynapse(neurons, "gi")
    gi.tau = 10.0
    gi.reversal = -20.0

def generate_connectivity():
    from ahh.util.np import JaggedAdjacencyMatrix
    global cm
    cm = JaggedAdjacencyMatrix(N)
    cm.connect_randomly(0.02, (0, N), (0, N))

def init_spike_transfer():
    from ahh.sim.connectivity import AtomicBuffer, AtomicReceiver, AtomicSend
    from ahh.sim import Conditional, ArrayLookup, StepAllocation

    global ge_buffer, ge_receiver
    global gi_buffer, gi_receiver
    global send
    global cm_alloc

    ge_buffer = AtomicBuffer.for_synapse(ge)
    ge_receiver = AtomicReceiver(neurons, ge, ge_buffer)
    ge_receiver.weight = 1.0

    gi_buffer = AtomicBuffer.for_synapse(gi)
    gi_receiver = AtomicReceiver(neurons, gi, gi_buffer)
    gi_receiver.weight = 1.0

    send = AtomicSend(neurons)
    send.target = Conditional("idx_model", 0, N, (
        (N_E, ge_buffer.outgoing),
        (N_I, gi_buffer.outgoing)
    ))
    flattened_cm = cm.flattened_with_rel_offsets
    cm_length = len(flattened_cm)
    cm_alloc = StepAllocation(sim, 'cm', cl_int, cm_length, True, 'cm')

    send.i_start = N
    send.i_max = ArrayLookup(cm_alloc, "idx_model")  # TODO
    send.neighbor_idx_lookup += ArrayLookup(cm_alloc, "2*idx_model + 1")
    send.neighbor_idx_lookup += N

def finalize():
    sim.finalize()

def allocate():
    sim.allocate()

def generate():
    sim.generate()

def print_cl_code():
    print sim.cl_code

def save_cl_code():
    from ahh.util import str_to_file
    str_to_file(sim.cl_code, "basic.cl")
    print "Wrote basic.cl."

def compile():
    sim.compile()



get_device()
init_sim()
init_neurons()
init_synapses()
#generate_connectivity()
#init_spike_transfer()
finalize()
generate()
save_cl_code()
